#ifndef MY_PARCSR_HPP
#define MY_PARCSR_HPP

#include <cstdio>
#include <cstdlib>
#include <cassert>
#include "mpi.h"
#include <vector>

/// @brief 并行CSR向量
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam calc_t 数据类型，可以为float或者double
template<typename idx_t, typename calc_t>
class MyParCSRVector
{
public:
    const MPI_Comm comm = MPI_COMM_NULL;// 通信域
    const idx_t glb_nrows;// 向量的全局行数
    const idx_t beg_row;// 本进程负责的起始行号
    const idx_t end_row;// 本进程负责的末尾行号（注意是开区间）
    calc_t * data = nullptr;// 本进程负责范围的数据
    MyParCSRVector(MPI_Comm _comm, const idx_t _glb_nrows, const idx_t _beg, const idx_t _end, const calc_t * dist_data):
        comm(_comm), glb_nrows(_glb_nrows), beg_row(_beg), end_row(_end)
    {
        const idx_t nrows = end_row - beg_row;
        data = new calc_t [nrows];
        if (dist_data) {
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++)
                data[i] = dist_data[i];
        } else {// 提前赋零以免出nan
            // int my_pid; MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
            // printf("P%d make zero %p\n", my_pid, &data);
            #pragma omp parallel for schedule(static)
            for (idx_t i = 0; i < nrows; i++)
                data[i] = 0.0;
        }
    }
    MyParCSRVector(const MyParCSRVector<idx_t, calc_t> & x):
        comm(x.comm), glb_nrows(x.glb_nrows), beg_row(x.beg_row), end_row(x.end_row)
    {
        const idx_t nrows = end_row - beg_row;
        data = new calc_t [nrows];
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < nrows; i++)
            data[i] = 0.0;
    }
    ~MyParCSRVector() {
        if (data) delete [] data;
    }
};

/// @brief 向量点积，内部以双精度进行计算
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 向量的数据类型，可以为float或者double
/// @param x 向量一
/// @param y 向量二
/// @return 内积结果 (x,y)
template<typename idx_t, typename data_t>
double vec_dot(const MyParCSRVector<idx_t, data_t> & x, const MyParCSRVector<idx_t, data_t> & y)
{
    assert(x.glb_nrows == y.glb_nrows && x.beg_row == y.beg_row && x.end_row == y.end_row);
    const idx_t nrows = x.end_row - x.beg_row;
    double loc_res = 0.0, glb_res;
    #pragma omp parallel for schedule(static) reduction(+:loc_res)
    for (idx_t i = 0; i < nrows; i++)
        loc_res += (double) x.data[i] * (double) y.data[i];
    MPI_Allreduce(&loc_res, &glb_res, 1, MPI_DOUBLE, MPI_SUM, x.comm);
    return glb_res;
}

/// @brief 向量范数，内部以双精度进行计算
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 向量的数据类型，可以为float或者double
/// @param x 向量
/// @return 向量范数 ||x||_2
template<typename idx_t, typename data_t>
double vec_norm(const MyParCSRVector<idx_t, data_t> & x)
{
    return sqrt(vec_dot(x, x));
}

/// @brief 向量加法
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 向量的数据类型，可以为float或者double
/// @param v1 源向量一
/// @param alpha 系数
/// @param v2 源向量二
/// @param v 目标向量
/// @result v = v1 + alpha * v2 
template<typename idx_t, typename data_t>
void vec_add(const MyParCSRVector<idx_t, data_t> & v1, data_t alpha,
            const MyParCSRVector<idx_t, data_t> & v2,
            MyParCSRVector<idx_t, data_t> & v)
{
    const idx_t nrows = v.end_row - v.beg_row;
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < nrows; i++)
        v.data[i] = v1.data[i] + alpha * v2.data[i];
}

/// @brief 向量加法
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 向量的数据类型，可以为float或者double
/// @param x 源向量
/// @param alpha 系数
/// @param y 目标向量
/// @result y = alpha * x
template<typename idx_t, typename data_t>
void vec_scale(const MyParCSRVector<idx_t, data_t> & x, data_t alpha,
            MyParCSRVector<idx_t, data_t> & y)
{
    const idx_t nrows = y.end_row - y.beg_row;
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < nrows; i++)
        y.data[i] = alpha * x.data[i];
}

/// @brief 向量拷贝（从源向量到目标向量的拷贝）
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t1 源向量的数据类型，可以为float或者double
/// @tparam data_t2 目标向量的数据类型，可以为float或者double
/// @param src 源向量
/// @param dst 目标向量
template<typename idx_t, typename data_t1, typename data_t2>
void vec_copy(const MyParCSRVector<idx_t, data_t1> & src, MyParCSRVector<idx_t, data_t2> & dst)
{
    assert(src.glb_nrows == dst.glb_nrows && src.beg_row == dst.beg_row && src.end_row == dst.end_row);
    const idx_t nrows = dst.end_row - dst.beg_row;
    #pragma omp parallel for schedule(static)
    for (idx_t i = 0; i < nrows; i++)
        dst.data[i] = src.data[i];
}

/// @brief 本地CSR矩阵
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 数据存储类型，可以为float或者double
/// @tparam calc_t 数据计算类型，可以为float或者double
template<typename idx_t, typename data_t, typename calc_t>
class MyCSRMatrix
{
public:
    const idx_t nrows;// 行数
    const idx_t ncols;// 列数
    const idx_t nnz;// 非零元数
    idx_t * row_ptr = nullptr;// 行偏移量
    idx_t * col_idx = nullptr;// 列索引
    data_t* vals = nullptr;// 非零元值

    /// @brief 本地CSR矩阵的构造函数，只开辟内存，不做初始化
    /// @param _nrows 行数
    /// @param _ncols 列数
    /// @param _nnz 非零元数
    MyCSRMatrix(const idx_t _nrows, const idx_t _ncols, const idx_t _nnz): nrows(_nrows), ncols(_ncols), nnz(_nnz) {
        row_ptr = new idx_t [nrows + 1];
        col_idx = new idx_t [nnz];
        vals    = new data_t[nnz];
    }
    ~MyCSRMatrix() {
        if (row_ptr) delete [] row_ptr;
        if (col_idx) delete [] col_idx;
        if (vals)    delete [] vals;
    }
    // y := alpha*A*x + beta*b
    void Mult(calc_t alpha, const calc_t *x, calc_t beta, const calc_t * b, calc_t *y) {
        #pragma omp parallel for schedule(static)
        for (idx_t i = 0; i < nrows; i++) {
            calc_t tmp = 0.0;
            for (idx_t p = row_ptr[i]; p < row_ptr[i+1]; p++) {
                tmp += vals[p] * x[col_idx[p]];
                // assert(tmp == tmp);
            }
            y[i] = alpha * tmp + beta * b[i];
        }
    }
    // y := alpha*A*x + beta*y
    void Mult(calc_t alpha, const calc_t * x, calc_t beta, calc_t *y) {
        Mult(alpha, x, beta, y, y);
    }
};

/// @brief 并行CSR矩阵
/// @tparam idx_t 索引类型，一般为int即够用
/// @tparam data_t 数据存储类型，可以为float或者double
/// @tparam calc_t 数据计算类型，可以为float或者double
template<typename idx_t, typename data_t, typename calc_t>
class MyParCSRMatrix
{
public:
    const MPI_Comm comm = MPI_COMM_NULL;// 通信域
    const idx_t glb_nrows;// 矩阵的全局行数
    const idx_t beg_row;// 本进程负责的起始行号
    const idx_t end_row;// 本进程负责的末尾行号（注意是开区间）
    idx_t * rows_partition = nullptr;// 记录各个进程的起始行号
    
    idx_t num_recvs = 0;// 有多少个邻居需要从他们那收取数据
    int   * recv_pids = nullptr;// 需要从哪些进程收？
    idx_t * recv_cnts = nullptr;// 从其它的每个进程处需要收多少个向量元素
    idx_t * col_map_offd = nullptr;// 从每个进程收的元素的行号，同时也记录了将全局坐标映射为非对角块格式的列序号

    idx_t num_sends = 0;// 有多少个邻居需要给他们发送自己负责的数据
    int   * send_pids = nullptr;// 需要给哪些进程发？
    idx_t * send_cnts = nullptr;// 给每个进程发的数量
    idx_t * send_row_ids = nullptr;// 给每个进程发的元素的行号

    MyCSRMatrix<idx_t, data_t, calc_t> * diag = nullptr;// 对角块矩阵
    MyCSRMatrix<idx_t, data_t, calc_t> * offd = nullptr;// 非对角块矩阵

    calc_t * send_buf = nullptr;// 发送缓冲区
    MPI_Request * send_reqs = nullptr;// 发送请求的句柄
    MPI_Request * recv_reqs = nullptr;// 接受请求的句柄

    /// @brief 并行CSR矩阵的构造函数
    /// @param _comm 通信域
    /// @param _glb_nrows 矩阵的全局行数 （隐含列数==行数）
    /// @param _beg 本进程负责的起始行号
    /// @param _end 本进程负责的末尾行号（注意是开区间）
    /// @param dist_row_ptr 分布式的行偏移
    /// @param dist_col_idx 分布式的列索引
    /// @param dist_vals 分布式的非零元值
    MyParCSRMatrix(MPI_Comm _comm, const idx_t _glb_nrows, const idx_t _beg, const idx_t _end,
        const idx_t * dist_row_ptr, const idx_t * dist_col_idx, const calc_t * dist_vals);
    ~MyParCSRMatrix();

    /// @brief 将自己负责区域、而又需被别人引用到的元素发出去，同时接收自己需要引用到别的进程的数据
    /// @param my_vec 本进程负责范围的向量，长度应为end_row - beg_row
    /// @param other_vec 其它进程负责范围的而本进程需要引用到的向量，长度应为offd->ncols
    void update_halo(const calc_t * my_vec, calc_t * other_vec);
    
    /// @brief 启动halo通信（不保证完成）
    /// @param my_vec 本进程负责范围的向量，长度应为end_row - beg_row
    /// @param other_vec 其它进程负责范围的而本进程需要引用到的向量，长度应为offd->ncols
    void init_update_halo(const calc_t * my_vec, calc_t * other_vec);

    /// @brief 等待上一次调用init_update_halo通信的完成
    void wait_update_halo();

    // y := alpha*A*x + beta*b
    void Mult(calc_t alpha, const MyParCSRVector<idx_t, calc_t> & x,
            calc_t beta, const MyParCSRVector<idx_t, calc_t> & b, MyParCSRVector<idx_t, calc_t> & y);
    // y := alpha*A*x + beta*y
    void Mult(calc_t alpha, const MyParCSRVector<idx_t, calc_t> & x,
            calc_t beta, MyParCSRVector<idx_t, calc_t> & y);
};
typedef MyParCSRMatrix<int, float , double> MyParCSRMatrix_I32F32F64;
typedef MyParCSRMatrix<int, double, double> MyParCSRMatrix_I32F64F64;
typedef MyParCSRMatrix<int, float , float > MyParCSRMatrix_I32F32F32;

#endif // MY_PARCSR_HPP